use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::mem;
use std::ptr::NonNull;
use std::alloc::{self, Layout};
use std::marker::PhantomData;
use anyhow::{Result, anyhow};
use serde::{Serialize, Deserialize};

/// 内存块句柄
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct MemoryHandle(u64);

impl MemoryHandle {
    pub fn new(id: u64) -> Self {
        Self(id)
    }
    
    pub fn as_u64(&self) -> u64 {
        self.0
    }
}

/// 内存块
struct MemoryBlock {
    ptr: NonNull<u8>,
    layout: Layout,
}

impl MemoryBlock {
    fn new(size: usize) -> Result<Self> {
        // 对齐到缓存行大小(64字节)
        let align = 64;
        let size = (size + align - 1) & !(align - 1);
        
        let layout = Layout::from_size_align(size, align)
            .map_err(|e| anyhow!("Invalid memory layout: {}", e))?;
            
        let ptr = unsafe {
            let ptr = alloc::alloc(layout);
            if ptr.is_null() {
                return Err(anyhow!("Memory allocation failed"));
            }
            NonNull::new_unchecked(ptr)
        };
        
        Ok(Self { ptr, layout })
    }
    
    fn as_slice(&self) -> &[u8] {
        unsafe {
            std::slice::from_raw_parts(self.ptr.as_ptr(), self.layout.size())
        }
    }
    
    fn as_slice_mut(&mut self) -> &mut [u8] {
        unsafe {
            std::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.layout.size())
        }
    }
}

impl Drop for MemoryBlock {
    fn drop(&mut self) {
        unsafe {
            alloc::dealloc(self.ptr.as_ptr(), self.layout);
        }
    }
}

/// 内存池，管理预分配的内存块
pub struct MemoryPool {
    // 按大小分类的可用内存块
    available_blocks: HashMap<usize, Vec<MemoryBlock>>,
    // 最大缓存块数量
    max_blocks_per_size: usize,
}

impl MemoryPool {
    pub fn new(max_blocks_per_size: usize) -> Self {
        Self {
            available_blocks: HashMap::new(),
            max_blocks_per_size,
        }
    }
    
    pub fn get_block(&mut self, size: usize) -> Result<MemoryBlock> {
        // 对齐到最近的2的幂
        let aligned_size = size.next_power_of_two();
        
        // 尝试从池中获取
        if let Some(blocks) = self.available_blocks.get_mut(&aligned_size) {
            if let Some(block) = blocks.pop() {
                return Ok(block);
            }
        }
        
        // 创建新块
        MemoryBlock::new(aligned_size)
    }
    
    pub fn return_block(&mut self, block: MemoryBlock) {
        let size = block.layout.size();
        
        let blocks = self.available_blocks
            .entry(size)
            .or_insert_with(Vec::new);
            
        // 只有当未达到最大缓存数量时才保留
        if blocks.len() < self.max_blocks_per_size {
            blocks.push(block);
        }
        // 否则让block在此作用域结束时被丢弃
    }
    
    pub fn cleanup(&mut self) {
        // 清理所有缓存的块
        self.available_blocks.clear();
    }
}

impl Default for MemoryPool {
    fn default() -> Self {
        Self::new(8) // 默认每个大小最多缓存8个块
    }
}

/// 共享内存管理器
pub struct SharedMemoryManager {
    next_handle: Arc<Mutex<u64>>,
    active_blocks: Arc<Mutex<HashMap<MemoryHandle, MemoryBlock>>>,
    memory_pool: Arc<Mutex<MemoryPool>>,
}

impl SharedMemoryManager {
    pub fn new() -> Self {
        Self {
            next_handle: Arc::new(Mutex::new(1)),
            active_blocks: Arc::new(Mutex::new(HashMap::new())),
            memory_pool: Arc::new(Mutex::new(MemoryPool::default())),
        }
    }
    
    pub fn allocate(&self, size: usize) -> Result<MemoryHandle> {
        // 获取新句柄
        let handle = {
            let mut next_handle = self.next_handle.lock().unwrap();
            let handle = MemoryHandle(*next_handle);
            *next_handle += 1;
            handle
        };
        
        // 分配内存块
        let block = {
            let mut pool = self.memory_pool.lock().unwrap();
            pool.get_block(size)?
        };
        
        // 存储活跃块
        let mut active_blocks = self.active_blocks.lock().unwrap();
        active_blocks.insert(handle, block);
        
        Ok(handle)
    }
    
    pub fn release(&self, handle: MemoryHandle) -> Result<()> {
        let mut active_blocks = self.active_blocks.lock().unwrap();
        
        if let Some(block) = active_blocks.remove(&handle) {
            let mut pool = self.memory_pool.lock().unwrap();
            pool.return_block(block);
            Ok(())
        } else {
            Err(anyhow!("Invalid memory handle: {:?}", handle))
        }
    }
    
    pub fn get_slice(&self, handle: MemoryHandle) -> Result<&[u8]> {
        let active_blocks = self.active_blocks.lock().unwrap();
        
        active_blocks.get(&handle)
            .map(|block| block.as_slice())
            .ok_or_else(|| anyhow!("Invalid memory handle: {:?}", handle))
    }
    
    pub fn get_slice_mut(&self, handle: MemoryHandle) -> Result<&mut [u8]> {
        let mut active_blocks = self.active_blocks.lock().unwrap();
        
        active_blocks.get_mut(&handle)
            .map(|block| block.as_slice_mut())
            .ok_or_else(|| anyhow!("Invalid memory handle: {:?}", handle))
    }
    
    pub fn cleanup(&self) {
        // 释放所有活跃块
        let mut active_blocks = self.active_blocks.lock().unwrap();
        active_blocks.clear();
        
        // 清理内存池
        let mut pool = self.memory_pool.lock().unwrap();
        pool.cleanup();
    }
}

impl Default for SharedMemoryManager {
    fn default() -> Self {
        Self::new()
    }
}

/// 内存视图，提供对共享内存区域的安全访问
pub struct MemoryView<'a, T: 'a> {
    manager: &'a SharedMemoryManager,
    handle: MemoryHandle,
    _marker: PhantomData<&'a T>,
}

impl<'a, T: 'a> MemoryView<'a, T> {
    pub fn new(manager: &'a SharedMemoryManager, size: usize) -> Result<Self> {
        let handle = manager.allocate(size)?;
        
        Ok(Self {
            manager,
            handle,
            _marker: PhantomData,
        })
    }
    
    pub fn handle(&self) -> MemoryHandle {
        self.handle
    }
    
    pub fn as_slice(&self) -> Result<&[u8]> {
        self.manager.get_slice(self.handle)
    }
    
    pub fn as_slice_mut(&self) -> Result<&mut [u8]> {
        self.manager.get_slice_mut(self.handle)
    }
}

impl<'a, T> Drop for MemoryView<'a, T> {
    fn drop(&mut self) {
        // 忽略错误，因为这是析构函数
        let _ = self.manager.release(self.handle);
    }
}